import numpy as np
import os
import json
from typing import Dict, Any, Union, List, Optional
from datetime import datetime


def save_counts_to_npy(counts: Dict[str, int], 
                       filepath: str, 
                       metadata: Optional[Dict[str, Any]] = None, 
                       to_probs: bool = False) -> str:
    """
    Save measurement counts dictionary to a NumPy file.
    
    Args:
        counts: Dictionary with measurement outcomes as keys and counts as values
        filepath: Path where the file should be saved (without extension)
        metadata: Optional dictionary with additional information to save
        to_probs: Whether to convert counts to probabilities

    Returns:
        Path to the saved file
    """
    # Create directory if it doesn't exist
    os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
    
    # Add .npy extension if not present
    if not filepath.endswith('.npy'):
        filepath += '.npy'
    
    # Prepare data dictionary with counts and metadata
    if to_probs:
        counts = {k: v / sum(counts.values()) for k, v in counts.items()}
    
    data_dict = {
        'counts': counts,
        'metadata': metadata or {},
        'timestamp': datetime.now().isoformat()
    }
    
    # Save to .npy file
    np.save(filepath, data_dict)
    
    return filepath


def load_counts_from_npy(filepath: str) -> Dict[str, Any]:
    """
    Load measurement counts and metadata from a NumPy file.
    
    Args:
        filepath: Path to the NumPy file
        
    Returns:
        Dictionary containing 'counts', 'metadata', and 'timestamp'
    """
    # Add .npy extension if not present
    if not filepath.endswith('.npy'):
        filepath += '.npy'
    
    # Load data
    data_dict = np.load(filepath, allow_pickle=True).item()
    
    return data_dict


def save_transition_matrix(transition_matrix: np.ndarray, 
                          basis_states: List[str], 
                          filepath: str,
                          qubit_indices: Optional[List[int]] = None,
                          backend_name: Optional[str] = None,
                          shots: Optional[int] = None) -> str:
    """
    Save a readout error transition matrix to a NumPy file.
    
    Args:
        transition_matrix: The readout error transition matrix
        basis_states: List of basis states corresponding to matrix rows/columns
        filepath: Path where the file should be saved (without extension)
        qubit_indices: Optional list of qubit indices used in characterization
        backend_name: Optional name of the backend used
        shots: Optional number of shots used for characterization
        
    Returns:
        Path to the saved file
    """
    # Create directory if it doesn't exist
    os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
    
    # Add .npy extension if not present
    if not filepath.endswith('.npy'):
        filepath += '.npy'
    
    # Prepare data dictionary
    data_dict = {
        'transition_matrix': transition_matrix,
        'basis_states': basis_states,
        'qubit_indices': qubit_indices,
        'backend_name': backend_name,
        'shots': shots,
        'timestamp': datetime.now().isoformat()
    }
    
    # Save to .npy file
    np.save(filepath, data_dict)
    

def load_transition_matrix(filepath: str) -> Dict[str, Any]:
    """
    Load a readout error transition matrix from a NumPy file.
    
    Args:
        filepath: Path to the NumPy file
        
    Returns:
        Dictionary containing the transition matrix and metadata
    """
    # Add .npy extension if not present
    if not filepath.endswith('.npy'):
        filepath += '.npy'
    
    # Load data
    data_dict = np.load(filepath, allow_pickle=True).item()
    
    return data_dict


def save_benchmark_results(benchmark_results: Dict[str, Any], 
                          filepath: str) -> str:
    """
    Save benchmark circuit execution results to a NumPy file.
    
    Args:
        benchmark_results: Dictionary containing benchmark results and metadata
        filepath: Path where the file should be saved (without extension)
        
    Returns:
        Path to the saved file
    """
    # Create directory if it doesn't exist
    os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
    
    # Add .npy extension if not present
    if not filepath.endswith('.npy'):
        filepath += '.npy'
    
    # Add timestamp if not present
    if 'timestamp' not in benchmark_results:
        benchmark_results['timestamp'] = datetime.now().isoformat()
    
    # Save to .npy file
    np.save(filepath, benchmark_results)
    
    return filepath


def load_benchmark_results(filepath: str) -> Dict[str, Any]:
    """
    Load benchmark results from a NumPy file.
    
    Args:
        filepath: Path to the NumPy file
        
    Returns:
        Dictionary containing benchmark results and metadata
    """
    # Add .npy extension if not present
    if not filepath.endswith('.npy'):
        filepath += '.npy'
    
    # Load data
    benchmark_results = np.load(filepath, allow_pickle=True).item()
    
    return benchmark_results


def merge_benchmark_results(filepaths: List[str], 
                          output_filepath: str) -> str:
    """
    Merge multiple benchmark result files into a single file.
    
    Args:
        filepaths: List of paths to benchmark result files
        output_filepath: Path where the merged file should be saved
        
    Returns:
        Path to the merged file
    """
    # Initialize merged results
    merged_results = {
        'merged_from': filepaths,
        'timestamp': datetime.now().isoformat(),
        'results': []
    }
    
    # Load and merge each file
    for filepath in filepaths:
        result = load_benchmark_results(filepath)
        merged_results['results'].append(result)
    
    # Save merged results
    return save_benchmark_results(merged_results, output_filepath)
